package com.whh.word2vec;

import org.deeplearning4j.models.embeddings.loader.WordVectorSerializer;
import org.deeplearning4j.models.word2vec.Word2Vec;
import org.deeplearning4j.text.sentenceiterator.BasicLineIterator;
import org.deeplearning4j.text.sentenceiterator.SentenceIterator;
import org.deeplearning4j.text.tokenization.tokenizer.preprocessor.CommonPreprocessor;
import org.deeplearning4j.text.tokenization.tokenizerfactory.DefaultTokenizerFactory;
import org.deeplearning4j.text.tokenization.tokenizerfactory.TokenizerFactory;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.io.IOException;

public class Test {
    private static Logger log = LoggerFactory.getLogger(Test.class);

    public static void main(String[] args) throws IOException {
        String filePath = "F:\\word2vec\\347.txt";
        log.info("Load & Vectorize Sentences....");
        SentenceIterator iter = new BasicLineIterator(new File(filePath));
        TokenizerFactory t = new DefaultTokenizerFactory();
        t.setTokenPreProcessor(new CommonPreprocessor());

        log.info("Building model....");
        Word2Vec vec = new Word2Vec.Builder().
                minWordFrequency(5).
                iterations(1)
                .layerSize(200)
                .minWordFrequency(2)
                .seed(42).
                        windowSize(5)
                .iterate(iter)
                .build();
        log.info("Fitting Word2Vec model....");
        vec.fit();
        log.info("Writing word vectors to text file....");

        // Write word vectors to file
        log.info("Writing word vectors to text file....");

        // Write word vectors to file
        WordVectorSerializer.writeWord2VecModel(vec, "F:\\word2vec\\pathToWriteto.txt");
        System.out.println("和微信最接近的10个词汇:" + vec.wordsNearest("人民检察院", 10));

    }
}
